# -*- coding:utf-8 -*-
# @Time    : 2023/5/5 12:47
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : planner.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import uuid

from bot.insurance_planner_gpt.utils.conversation_process_util import ConversationProcessUtil
from bot.insurance_planner_gpt.utils.seat_planner_msg_channel import SeatPlannerMsgChannel
from bot.insurance_planner_gpt.utils.causal_data_process_util import CausalDataProcessUtil
from common.message import ConversationMessage, Message, MessageContent
from bot.insurance_planner_gpt.agent.conversation_agent import Conversation
from bot.insurance_planner_gpt.agent.user_questioning import QuestioningDisputeResolution
from bot.insurance_planner_gpt.agent.user_info_extract import UserInfoExtract
from bot.insurance_planner_gpt.agent.dynamic_matching import MatchSolution
from bot.insurance_planner_gpt.user_info import UserInfo
from db.redis.redis_service import RedisMemory
from bot.insurance_planner_gpt.utils.user_info_util import ProfessCate
import json
from common import constants
import asyncio
from common.log import logger
from bot.bot import Bot
from util_tool import utils

memory_cache = RedisMemory()
process_seat_channel = SeatPlannerMsgChannel()
profess = ProfessCate()
# condipro = CandidateProducts()
condipro = CausalDataProcessUtil.candidate_product

class PlannerGPT(Bot):

    def __init__(self, client_type, client_ws=None, planner_role="", company_name="", company_business="",
                 company_values=""):
        self.professional_knowledge = []
        self.client_ws = client_ws
        self.client_type = client_type
        self.solution_info = ""
        self.solution_other_messages = []
        self.user_question = ""
        self.background_knowledge = ""
        self.user_info_cn = {}
        self.user_info_update = False
        self.conversation_history_str = ""
        self.next_solution = ""
        self.insurance_coverage_str = ""
        self.planner_role = planner_role
        self.company_name = company_name
        self.company_business = company_business
        self.company_values = company_values

    """Controller model for the planner Agent."""
    user_info_dict = UserInfo.user_info_dict

    # knowledge_base_query = BGEQuery()
    # knowledge_base_query.load_all_knowledge_base()

    # async def calculate_insurance_coverage(self):
    #     self.solution_info = await PremiumCalculation(self.user_info_en).calculate_insurance_coverage()

    # 提取用户信息
    async def extract_user_info(self, session_id):

        user_info_rule = UserInfo.user_info_rule
        user_info_extract = UserInfoExtract(role=self.planner_role,
                                            conversation_history=self.conversation_history_str,
                                            user_info_rule=user_info_rule)

        user_info_str = await user_info_extract.achat_auto_llm(type="self")

        logger.info(f"会话id:{self.session_id},获取的用户信息是:{user_info_str}")
        try:
            user_info_cn = utils.json_load_llm_result(user_info_str)
        except Exception as e:
            logger.exception(e)
            user_info_cn = {}
        # user_info_cache = {}
        # self.user_info_cn = {}
        self.user_info_cn = user_info_cn
        self.candidate_products = []

        # if memory_cache.exist_key("planner:task:" + session_id):
        #     user_info_cache = memory_cache.get_memory("planner:task:" + session_id, "UserInfoExtract")
        #     if user_info_cache:
        #         self.user_info_cn = user_info_cache

        # 判断user_info_cn是否所有key和value都在user_info_cache中，是的话
        # if not utils.is_dict_in_dict(user_info_cn, user_info_cache):
        #     self.user_info_update = True
        #     self.user_info_cn.update(user_info_cn)
        try:
            if "当前关注保险产品" in self.user_info_cn:
                all_product_name = self.user_info_cn["当前关注保险产品"]
                if isinstance(all_product_name, list):
                    all_product_name = ','.join(all_product_name)
                all_product_name = all_product_name.strip().split(',')
                if len(all_product_name) > 0:
                    # product_name_lst = []
                    for product_name in all_product_name:
                        candidate_prod = condipro.get_condidate_products(target_product_query=product_name)
                        if candidate_prod:
                            self.candidate_products.extend(candidate_prod)


        except Exception as e:
            logger.exception(e)

        memory_cache.set_memory("planner:task:" + session_id, "UserInfoExtract", self.user_info_cn)

        # 发送给坐席平台
        user_info_msg = {'title': '用户信息：', 'content': self.user_info_cn}
        msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_AGENT_SOLUTION, session_id, self.client_type,
                                          json.dumps(user_info_msg, ensure_ascii=False),
                                          session_id, role=constants.ROLE_ASSISTANT,
                                          message_id=str(uuid.uuid4()).replace("-", ""), action="agent_solution")
        await process_seat_channel.publish_msg(process_seat_channel.planner_channel, msg)

        # 根据user_info_dict把user_info_en中的英文key替换成中文key
        self.user_info_en = UserInfo.get_english_dict(self.user_info_cn)

    async def get_product_causal_solution(self, plan_knowledge_infos):
        # 确定产品类型和产品名称
        product_solution = MatchSolution()
        result = await product_solution.causal_match(self.session_id, self.user_question, self.user_info_cn, self.candidate_products,
                                                     plan_knowledge_infos,
                                                     self.conversation_history_str)
        return result

    async def count_conversation_rounds(self):
        round_count = 0
        try:
            user_spoke = False
            assistant_spoke = False
            for entry in self.context_list:
                if entry['role'] == 'user':
                    user_spoke = True
                elif entry['role'] == 'assistant':
                    assistant_spoke = True
                # Check if both user and assistant have spoken in this round
                if user_spoke and assistant_spoke:
                    round_count += 1
                    user_spoke = False  # Reset for the next round
                    assistant_spoke = False  # Reset for the next round
        except Exception as e:
            logger.error(f"【统计对话轮次异常，返回0轮】：{e}")
        return round_count


    async def generate_reply(self, is_stream=False):
        if self.solution_info:
            memory_cache.set_solution_memory(self.session_id, self.solution_info)
        # 话术生成
        utterance_bot = Conversation(self.planner_role, self.company_name, self.company_business,
                                     self.company_values, self.conversation_history_str,
                                     self.next_solution,
                                     str(self.solution_info), self.user_info_cn)
        # 前3轮用自有模型，后面的用gpt
        conversation_rounds = await self.count_conversation_rounds()
        model_type = "self" if  conversation_rounds<= 3 else "gpt"
        logger.info(f"会话id:{self.session_id},【对话轮次】：{conversation_rounds}, 话术调的模型:{model_type}")
        if is_stream:
            replay_text = await utterance_bot.achat_auto_llm(type=model_type, is_stream=is_stream)
        else:
            # 【仅在非流式输出时检查生成问题】
            #  检查生成的连接和数字是否都是prompt里的知识，如果不是则重新生成话术，最多重试3次
            all_included = False
            retry_times = 0
            while not all_included and retry_times < 3:
                replay_text = await utterance_bot.achat_auto_llm(type=model_type, is_stream=is_stream)
                all_included = await utterance_bot.check_reply_inclusion(replay_text)
                retry_times = retry_times + 1
            logger.info(f"会话id:{self.session_id},【话术生成次数】：{retry_times}")
        return replay_text

    async def get_user_question(self, session_id):
        try:
            # 历史用户问题查询
            user_question_cache = ""
            if memory_cache.exist_key("planner:task:" + session_id):
                user_question_cache = memory_cache.get_memory("planner:task:" + session_id,
                                                              "QuestioningDisputeResolution")
            questioning_dispute_resolution = QuestioningDisputeResolution(self.planner_role,
                                                                          self.conversation_history_str,
                                                                          user_question_cache)
            self.user_question = await questioning_dispute_resolution.achat_auto_llm(type="self")
            # 更新用户问题
            memory_cache.set_memory("planner:task:" + session_id, "QuestioningDisputeResolution", self.user_question)
            # 发送给坐席
            user_info_msg = {'title': '用户问题：', 'content': self.user_question}
            msg = ConversationMessage.gen_msg(constants.MESSAGE_TYPE_AGENT_SOLUTION, session_id, self.client_type,
                                              json.dumps(user_info_msg, ensure_ascii=False),
                                              session_id, role=constants.ROLE_ASSISTANT,
                                              message_id=str(uuid.uuid4()).replace("-", ""), action="agent_solution")
            await process_seat_channel.publish_msg(process_seat_channel.planner_channel, msg)

            logger.info(f"会话id：{self.session_id},整理后的用户问题是:{self.user_question}")
        except Exception as e:
            # 打印错误的堆栈信息
            logger.exception(e)
            self.user_question = ""

    async def async_reply(self, context, session_id, client_type, channel_need_transfer=False, is_stream=False):
        conversation = ConversationProcessUtil.format_conversation_history(context,
                                                                           assistant_role_name=self.planner_role)
        self.session_id = session_id
        self.client_type = client_type
        self.solution_other_messages = []
        self.context_list = context
        self.conversation_history_str = conversation

        if len(self.context_list) > 1:
            # 确定用户问题 和 用户信息
            await asyncio.gather(self.extract_user_info(session_id), self.get_user_question(session_id))
            # quesitons = self.user_question.split("\n")
            plan_knowledge_infos = []
            # 动态mapping（用户需求->方案匹配->沟通策略）
            result = await self.get_product_causal_solution(plan_knowledge_infos)
            logger.info(f"会话id:{session_id},动态mapping结果:{result}")

            product_solution_cots = result["匹配方案"]
            communication_strategy_str = result["沟通策略"]
            # plan_description = result["方案说明"]
            background_knowledge = result["条款知识"]
            insurance_coverage = result["保费计算"]
            company_knowledge = result["公司知识"]

            self.user_need_info_list = []
            if background_knowledge:
                self.solution_info = self.solution_info + '【条款知识】:' + background_knowledge + '\n\n'
            if insurance_coverage and insurance_coverage != "{}":
                self.solution_info = self.solution_info + '【保额保费预估】：' + insurance_coverage + '\n\n'
            if company_knowledge:
                self.solution_info = self.solution_info +'【'+ self.company_name +'-公司知识】:'+ company_knowledge + '\n\n'

            # 更新潜在需求预测数据
            self.next_solution = communication_strategy_str
            logger.info(f"会话id:{session_id},整合后的下一步策略为:{self.next_solution},")

        else:
            self.next_solution = "简介：介绍自己用热情活泼的语气保持轻松谈话的氛围。问候语应该是热情友好的。询问用户有什么可以帮助他的。"
        result = await self.generate_reply(is_stream)


        # 判断是否转人
        # if channel_need_transfer:
        #     transfer_res = await planner_judge_transfer(self.conversation_history_str)
        #     plan_assume_str = {'title': 'transfer_agent：', 'content': transfer_res}
        #     self.solution_other_messages.extend(plan_assume_str)

        return result, self.solution_other_messages

    def post_process_code(self, code):
        sep = "\n```"
        if sep in code:
            blocks = code.split(sep)
            if len(blocks) % 2 == 1:
                for i in range(1, len(blocks), 2):
                    blocks[i] = blocks[i].replace("\\_", "_")
            code = sep.join(blocks)
        return code


class PlannerChat:

    def __init__(self, client_ws, planner_role="", company_name="", company_business="",
                 company_values=""):
        self.client_type = None
        self.client_ws = client_ws
        self.planner_role = planner_role
        self.company_name = company_name
        self.company_business = company_business
        self.company_values = company_values
        self.planner = PlannerGPT(self.client_type, self.client_ws, self.planner_role, self.company_name,
                                  self.company_business,
                                  self.company_values)

    async def async_reply(self, context, session_id, client_type, channel_need_transfer=False, is_stream=False):
        self.client_type = client_type
        result, other_messages = await self.planner.async_reply(context, session_id, client_type, channel_need_transfer,
                                                                is_stream)
        return result, other_messages

    async def get_agent_reply(self, context, session_id, client_type, channel_need_transfer=False, is_stream=False):
        self.client_type = client_type
        result, other_messages = await self.planner.async_reply(context, session_id, client_type, channel_need_transfer,
                                                                is_stream)
        return result, other_messages

    def reply(sefl, query, user_id):
        return ''

    async def send_msg(self, msg_type, client_type, session_id, message, code, message_id, target_client_id="",
                       client_id="", role=constants.ROLE_ASSISTANT):
        if self.client_ws:
            # 构造消息格式
            send_content = MessageContent(msg_type, message)
            content_arr = [send_content]
            send_message = Message(message_id, role, content_arr, None)
            send_msg = ConversationMessage(client_type, client_id, session_id, '', target_client_id, send_message, code)
            await self.client_ws.send(send_msg.to_json())
        else:
            print("客户端 WebSocket 连接未建立")

    async def send_multi_msg(self, client_type, session_id, multi_message, code):
        if self.client_ws:
            # 构造多条消息格式
            content_arr = []
            for message in multi_message:
                if isinstance(message, MessageContent):
                    send_content = MessageContent(message.msg_type, message.data)
                    content_arr.append(send_content)
            send_message = Message(str(uuid.uuid4()).replace('-', ''), constants.ROLE_ASSISTANT, content_arr, None)
            send_msg = ConversationMessage(client_type, session_id, '', send_message, code)

            await self.client_ws.send(send_msg.to_json())
        else:
            print("客户端 WebSocket 连接未建立")


if __name__ == '__main__':
    planner_bot = PlannerGPT(None)
    context = [["assistant", "您好！我是一名退休规划师，真诚、温暖和专业是我们公司的价值观。我们公司为用户提供全面的退休服务，以帮助您实现个性化的退休策略。请问您现在有什么可以帮助我们帮您的吗？"],
               ["user", "我今年35岁 性别男 退休后每月需要花费8000元,30岁缴纳的社保 另外还有两年的灵活就业社保"]
               ]
    # 提取用户信息测试

    asyncio.run(planner_bot.async_reply(context, "", ""))
